rm(list = ls())
library(patchwork)
library(tidyverse)
library(ggplot2)
library(matlab)
# parameters setting
psi <- 0.05
beta1 <- 50
v <- 0.2
theta <- 0.5
mu <- 0.01

# exo variables
mt <- 100
beta0 <- 2100
ytn <- 2000

#
ode_solve <- function(psi, beta1,v, theta, mu, mt, beta0,ytn){
  # coef matrix
  A <- matrix(c(0,mu,-v*beta1/theta,v*(beta1*mu-beta1*psi/theta-1)),2,byrow = T)
  B <- matrix(c(0,0,-mu,v,v*beta1/theta,-v*beta1*mu),2,byrow =T)
  zt <- matrix(c(beta0,mt,ytn),ncol = 1)

  # steady state
  st <- -solve(A) %*% B %*% zt
  # eig value
  ev <- eigen(A)$values
  return(list(A=A, B=B, zt=zt, st=st, ev=ev))
}

rlt <- ode_solve(psi, beta1,v, theta, mu, mt, beta0,ytn)
rlt

# IRF: m0
mt <- 101
new_rlt <- ode_solve(psi, beta1,v, theta, mu, mt, beta0,ytn)

x <- matrix(0,nrow = 4, ncol = 20, dimnames = list(c('p','y','yd','i')))
for (i in 1:20) {
  if ( i == 1){
    x[1:2,1] <- rlt$st
  }else {
    x[1:2,i] <- (new_rlt$A + eye(2)) %*% x[1:2,i-1] + new_rlt$B %*% new_rlt$zt # xt = (A+I)x_t-1+ Bz
  }
  x[4,i] <- -(mt-x[1,i]-psi*x[2,i])/theta
  x[3,i] <- beta0-beta1*(x[4,i]-mu*(x[2,i]-ytn))
}
picdata <- t(x) %>% as.data.frame()
p1 <- ggplot(picdata, aes(x = 1:nrow(picdata),y = y)) + geom_line() + labs(x='') +theme_bw()
p2 <- ggplot(picdata, aes(x = 1:nrow(picdata),y = p)) + geom_line() + labs(x='')+ theme_bw()
p3 <- ggplot(picdata, aes(x = 1:nrow(picdata),y = yd)) + geom_line() + labs(x='')+ theme_bw()
p4 <- ggplot(picdata, aes(x = 1:nrow(picdata),y = i)) + geom_line() + labs(x='')+ theme_bw()
(p1 + p3)/(p2+p4)
ggsave('../m_irf.pdf')

# sensitivity:
psi <-  0.01
sen_rlt <- ode_solve(psi, beta1,v, theta, mu, mt, beta0,ytn)

x <- matrix(0,nrow = 4, ncol = 30, dimnames = list(c('p','y','yd','i')))
for (i in 1:30) {
  if ( i == 1){
    x[1:2,1] <- rlt$st
  }else {
    x[1:2,i] <- (sen_rlt$A + eye(2)) %*% x[1:2,i-1] + sen_rlt$B %*% new_rlt$zt # xt = (A+I)x_t-1+ Bz
  }
  x[4,i] <- -(mt-x[1,i]-psi*x[2,i])/theta
  x[3,i] <- beta0-beta1*(x[4,i]-mu*(x[2,i]-ytn))
}
picdata <- t(x) %>% as.data.frame()
p1 <- ggplot(picdata, aes(x = 1:nrow(picdata),y = y)) + geom_line() + labs(x='') +theme_bw()
p2 <- ggplot(picdata, aes(x = 1:nrow(picdata),y = p)) + geom_line() + labs(x='')+ theme_bw()
p3 <- ggplot(picdata, aes(x = 1:nrow(picdata),y = yd)) + geom_line() + labs(x='')+ theme_bw()
p4 <- ggplot(picdata, aes(x = 1:nrow(picdata),y = i)) + geom_line() + labs(x='')+ theme_bw()
(p1 + p3)/(p2+p4)
ggsave('../m_sen.pdf')
